import os.path
import cv2
import insightface
import numpy as np

from sklearn import preprocessing
from util.http import fail_api, success_api


class FaceRecognition():

    def __init__(self, gpu_id = 0, face_db = r'C:\Users\admin\Desktop\project\facerecognize\face_db', threshold = 1.0, det_thresh = 0.50, det_size = (640, 640)):
        self.gpu_id = gpu_id
        self.face_db = face_db
        self.threshold = threshold
        self.det_thresh = det_thresh
        self.det_size = det_size

        #加载人脸分析模型
        self.model = insightface.app.FaceAnalysis(root='./', allowed_modules=['detection', 'recognition'])
        #配置insightface模型参数
        self.model.prepare(ctx_id = self.gpu_id, det_thresh = self.det_thresh, det_size = self.det_size)
        #存储人脸库的人脸特征
        self.faces_embedding = list()
        #加载人脸库中的人脸
        self.load_faces(self.face_db)

    #加载人脸库
    def load_faces(self, face_db_path):
        if not os.path.exists(face_db_path):
            os.makedirs(face_db_path)
        for root, dirs, files in os.walk(face_db_path):
            for file in files:
                # 从文件路径读取图像并解码为 OpenCV 可处理的格式，尤其适用于处理中文路径或特殊字符路径的情况
                #os.path.join(root, file)：跨平台拼接目录和文件名（如 "data/图片.jpg"）
                #np.fromfile(..., dtype=np.uint8)：将文件读取为 uint8 格式的 NumPy 数组（二进制流）
                #cv2.imdecode(..., 1)：将二进制流解码为 OpenCV 的 BGR 格式图像
                path = os.path.join(root, file)
                img_uint8 =  np.fromfile(path, dtype=np.uint8)
                input_image = cv2.imdecode(img_uint8, 1)
                if input_image is not None:
                    faces = self.model.get(input_image)
                    if faces:
                        face =faces[0]
                        embedding = np.array(face.embedding).reshape((1, -1))  # 转为1x512的NumPy数组
                        # 特征向量归一化（单位向量化） 归一化：消除特征向量的模长影响，使相似度计算（如余弦相似度）仅依赖方向差异。
                        embedding = preprocessing.normalize(embedding)
                        self.faces_embedding.append({
                            "user_name":file.split(".")[0],
                            "feature": embedding
                        })
                        print(file.split(".")[0])
                    else:
                        print(f"未在{file}中检测到人脸")

    def load_faces_sim(self, face_db_path):
        if not os.path.exists(face_db_path):
            os.makedirs(face_db_path)
        for root, dirs, files in os.walk(face_db_path):
            for file in files:
                # 从文件路径读取图像并解码为 OpenCV 可处理的格式，尤其适用于处理中文路径或特殊字符路径的情况
                #os.path.join(root, file)：跨平台拼接目录和文件名（如 "data/图片.jpg"）
                #np.fromfile(..., dtype=np.uint8)：将文件读取为 uint8 格式的 NumPy 数组（二进制流）
                #cv2.imdecode(..., 1)：将二进制流解码为 OpenCV 的 BGR 格式图像
                path = os.path.join(root, file)
                img_uint8 =  np.fromfile(path, dtype=np.uint8)
                input_image = cv2.imdecode(img_uint8, 1)
                if input_image is not None:
                    faces = self.model.get(input_image)
                    if faces:
                        face =faces[0]
                        embedding = np.array(face.embedding).reshape((1, -1))  # 转为1x512的NumPy数组
                        self.faces_embedding.append({
                            "user_name":file.split(".")[0],
                            "feature": embedding
                        })
                        print(file.split(".")[0])
                    else:
                        print(f"未在{file}中检测到人脸")

    def register(self, image, user_name):
        faces = self.model.get(image)
        if len(faces) != 1:
            return fail_api("图片检测不到人脸")
        embedding = np.array(faces[0].embedding).reshape((1, -1))
        embedding = preprocessing.normalize(embedding)
        is_exits = False
        for com_face in self.faces_embedding:
            r = self.feature_compare(embedding, com_face["feature"], self.threshold)
            if r:
                is_exits = True
        if is_exits:
            response = fail_api("该用户已存在")
            return response
        # cv2.imdecode('.png', image)[1].tofile(os.path.join(self.face_db, '%s.png'%user_name))
        if not isinstance(image, np.ndarray):
            return fail_api("输入必须是 NumPy 数组")
        success, encoded_image = cv2.imencode('.png', image)
        if not success:
            return fail_api("图像编码失败")
        file_path = os.path.join(self.face_db, f'{user_name}.png')
        with open(file_path, 'wb') as f:
            f.write(encoded_image.tobytes())
        self.faces_embedding.append({
            "user_name": user_name,
            "feature": embedding
        })
        return "success"

    def register_sim(self, image, user_name):
        faces = self.model.get(image)
        if len(faces) != 1:
            return fail_api("图片检测不到人脸")
        embedding = np.array(faces[0].embedding).reshape((1, -1))
        is_exits = False
        for com_face in self.faces_embedding:
            r = self.feature_compare_sim(embedding, com_face["feature"])
            if r:
                is_exits = True
        if is_exits:
            response = fail_api("该用户已存在")
            return response
        # cv2.imdecode('.png', image)[1].tofile(os.path.join(self.face_db, '%s.png'%user_name))
        if not isinstance(image, np.ndarray):
            return fail_api("输入必须是 NumPy 数组")
        success, encoded_image = cv2.imencode('.png', image)
        if not success:
            return fail_api("图像编码失败")
        file_path = os.path.join(self.face_db, f'{user_name}.png')
        with open(file_path, 'wb') as f:
            f.write(encoded_image.tobytes())
        self.faces_embedding.append({
            "user_name": user_name,
            "feature": embedding
        })
        return "success"

    def recognition(self, image):
        faces = self.model.get(image)
        if len(faces) != 1:
            return fail_api("上传的图片检测不到人脸")
        embedding = np.array(faces[0].embedding).reshape((1, -1))
        embedding = preprocessing.normalize(embedding)
        user_name = '未识别_'
        results = set()
        for com_face in self.faces_embedding:
            r = self.feature_compare(embedding, com_face["feature"], self.threshold)
            if r:
                user_name = com_face["user_name"]
                results.add(user_name)
        return results

    def recognition_sim(self, image):
        faces = self.model.get(image)
        if len(faces) != 1:
            return fail_api("上传的图片检测不到人脸")
        embedding = np.array(faces[0].embedding).reshape((1, -1))
        user_name = '未识别_'
        results = set()
        for com_face in self.faces_embedding:
            r = self.feature_compare_sim(embedding, com_face["feature"], self.threshold)
            if r:
                user_name = com_face["user_name"]
                results.add(user_name)
        return results
    #欧式距离
    @staticmethod
    def feature_compare(feature1, feature2, threshold):
        diff = np.subtract(feature1, feature2)
        dist = np.sum(np.square(diff), 1)
        if dist < threshold:
            return True
        else:
            return False

    #余弦相似度
    @staticmethod
    def feature_compare_sim(feature1, feature2, threshold = 0.8):
        # 归一化（可选但推荐）
        feature1 = feature1 / np.linalg.norm(feature1)
        feature2 = feature2 / np.linalg.norm(feature2)
        cosine_sim = np.dot(feature1, feature2.T)
        if cosine_sim >= threshold:
            return True
        else:
            return False


fr = FaceRecognition()